import os
import shutil
from matplotlib.font_manager import json_load
import numpy as np
import cv2

yolo_train_img_path = "./yolo/train/images/"
yolo_val_img_path = "./yolo/val/images/"
yolo_train_label_path = "./yolo/train/labels/"
yolo_val_label_path = "./yolo/val/labels/"
os.makedirs(yolo_train_img_path, exist_ok=True)
os.makedirs(yolo_val_img_path, exist_ok=True)
os.makedirs(yolo_train_label_path, exist_ok=True)
os.makedirs(yolo_val_label_path, exist_ok=True)

labelme_data_path = "C:/Users/kang_/Desktop/data"
idx = 0
for root, dirs, files in os.walk(labelme_data_path):
    for file in files:
        if not file.endswith(".json"):
            continue
        label_path = os.path.join(root, "1280x720_MA.json")
        image_path = os.path.join(root, "1280x720_MA.png")
        label_json = json_load(label_path)
        cls_idx, x, y, w, h = 0, 0, 0, 0, 0
        pn = []
        for shape in label_json["shapes"]:
            label = shape["label"]
            points = shape["points"]
            if label == "pig":
                pig_segment = np.array(points, dtype=np.float32)
                x, y, w, h = cv2.boundingRect(pig_segment)

                x_center = ((x + w / 2) + 10) / 1280.0
                y_center = ((y + h / 2) + 10) / 780.0
                w /= 1280.0
                h /= 780.0

            if label == "pig_pose":
                points = shape["points"]
                pn.append(points[0])

        if len(pn) == 0:
            continue
        yolo = [cls_idx, x_center, y_center, w, h]
        for i in pn:
            yolo.append(i[0] / 1280.0)
            yolo.append(i[1] / 780.0)
        lable = " ".join([str(i) for i in yolo])
        dir = root.split("\\")
        dir = dir[1] + "_" + dir[2]
        if idx <= 100:
            with open(yolo_val_label_path + f"{dir}.txt", "w") as fp:
                fp.write(lable)
            shutil.copy(image_path, yolo_val_img_path + f"{dir}.png")
        else:
            with open(yolo_train_label_path + f"{dir}.txt", "w") as fp:
                fp.write(lable)
            shutil.copy(image_path, yolo_train_img_path + f"{dir}.png")
        idx += 1
